/**
 * 
 */
package org.apache.lucene.search;

import java.io.IOException;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.grouping.GroupDocs;
import org.apache.lucene.search.grouping.SearchGroup;
import org.apache.lucene.search.grouping.TermSecondPassGroupingCollector;
import org.apache.lucene.search.grouping.TopGroups;

/**
 * @author Liuye
 * 
 */
public class B2BTermSecondPassGroupingCollector extends
		TermSecondPassGroupingCollector {
	private final int maxDocsPerGroup;
	protected SearchGroupDocs<String>[] groupDocs;
	private final Collection<SearchGroup<String>> groups;
	private final Sort withinGroupSort;
	private final Sort groupSort;

	private int totalHitCount;
	private int totalGroupedHitCount;

	private String groupField;
	private Map<String, Integer> docsPerGroupCounter;
	private FieldCache.StringIndex index;

	public B2BTermSecondPassGroupingCollector(String groupField,
			Collection<SearchGroup<String>> groups, Sort groupSort,
			Sort withinGroupSort, int maxDocsPerGroup, boolean getScores,
			boolean getMaxScores, boolean fillSortFields) throws IOException {
		super(groupField, groups, groupSort, withinGroupSort, maxDocsPerGroup,
				getScores, getMaxScores, fillSortFields);
		// System.out.println("SP init");
		if (groups.size() == 0) {
			throw new IllegalArgumentException(
					"no groups to collect (groups.size() is 0)");
		}
		this.groupField = groupField;
		this.groupSort = groupSort;
		this.withinGroupSort = withinGroupSort;
		this.groups = groups;
		this.maxDocsPerGroup = maxDocsPerGroup;
		this.docsPerGroupCounter = new HashMap<String, Integer>();
	}

	@Override
	public void collect(int doc) throws IOException {
		String groupValue = getDocGroupValue(doc);
		Integer docsPerGroup = docsPerGroupCounter.get(groupValue);
		if (docsPerGroup == null || docsPerGroup < maxDocsPerGroup) {
			docsPerGroup = docsPerGroup == null ? 1 : ++docsPerGroup;
			docsPerGroupCounter.put(groupValue, docsPerGroup);
			totalHitCount++;
		}
		
		SearchGroupDocs<String> group = retrieveGroup(doc);
		if (group != null) {
			totalGroupedHitCount++;
			group.collector.collect(doc);
		}
	}

	@Override
	public void setNextReader(IndexReader reader, int docBase)
			throws IOException {
		super.setNextReader(reader, docBase);
		index = FieldCache.DEFAULT.getStringIndex(reader, groupField);
	}

	@Override
	public TopGroups<String> getTopGroups(int withinGroupOffset) {
		@SuppressWarnings("unchecked")
		final GroupDocs<String>[] groupDocsResult = (GroupDocs<String>[]) new GroupDocs[groups
				.size()];
		int groupIDX = 0;
		for (SearchGroup<String> group : groups) {
			final SearchGroupDocs<String> groupDocs = groupMap
					.get(group.groupValue);
			final TopDocs topDocs = groupDocs.collector.topDocs(
					withinGroupOffset, maxDocsPerGroup);
			groupDocsResult[groupIDX++] = new GroupDocs<String>(
					topDocs.getMaxScore(), topDocs.totalHits,
					topDocs.scoreDocs, groupDocs.groupValue, group.sortValues);
		}

		return new TopGroups<String>(groupSort.getSort(),
				withinGroupSort == null ? null : withinGroupSort.getSort(),
						totalHitCount, totalGroupedHitCount, groupDocsResult);
	}

	protected String getDocGroupValue(int doc) {
		final int ord = index.order[doc];
		return ord == 0 ? null : index.lookup[ord];
	}

	protected int calcTotalHitCount() {
		int tatalHitCount = 0;
		for (Map.Entry<String, Integer> entry : docsPerGroupCounter
				.entrySet()) {
			tatalHitCount += entry.getValue();
		}
		return tatalHitCount;
	}
}
